#include "com.h"

using namespace std;

CoMUtilities::CoMUtilities(double my_support_polygon_x, double my_support_polygon_y, bool visualize_com_in_pcl)
{
  support_polygon_x = my_support_polygon_x;
  support_polygon_y = my_support_polygon_y;

  listener = new tf::TransformListener(ros::Duration(180)); // it must be scoped to persist because of its cache

  init_links_params();
  init_robovie_x_mass();

  pcl_com_visualization = visualize_com_in_pcl;

  if (pcl_com_visualization)
  {
    p = new pcl::visualization::PCLVisualizer ("PCL Viewer");
    p->setBackgroundColor(0, 0, 0);
    p->addCoordinateSystem(0.01);
  }
}

void CoMUtilities::init_links_params() {
	
	vector<double> v(4);

    v[0] = 0.4265;
    v[1] = 0;
    v[2] = 0;
    v[3] = 0.241;
	tfs["torso"] = v;
    
    v[0] = 0.06;
    v[1] = 0;
    v[2] = 0;
    v[3] = 0.02925;
	tfs["head"] = v;

    v[0] = 0.007;
    v[1] = 0;
    v[2] = 0;
    v[3] = 0;
	tfs["right_shoulder"] = v;
	
    v[0] = 0.007;
    v[1] = 0;
    v[2] = 0;
    v[3] = 0;
	tfs["left_shoulder"] = v;

    v[0] = 0.0885;
    v[1] = 0.0;
    v[2] = -0.002279;
    v[3] = -0.029454;
	tfs["right_arm"] = v;

    v[0] = 0.0885;
    v[1] =  0.0;
    v[2] = 0.002279;
    v[3] =-0.029454;
	tfs["left_arm"] = v;
	
    v[0] = 0.021;
    v[1] = 0.0;
    v[2] = 0.0;
    v[3] = -0.041499;
	tfs["right_hand"] = v;
	
    v[0] = 0.021;
    v[1] = 0.0;
    v[2] = -0.0;
    v[3] =-0.041499;
	tfs["left_hand"] = v;

    v[0] = 0.013; //0.1;
    v[1] = 0.0;
    v[2] = 0.0;
    v[3] = -0.016005;
	tfs["right_hip"] = v;

    v[0] = 0.013; //0.1;
    v[1] = 0.0;
    v[2] = -0.0;
    v[3] = -0.016005;
	tfs["left_hip"] = v;

    v[0] = 0.0885; //0.1;
    v[1] = 0.00228;
    v[2] = 0.0; 
    v[3] =-0.02945;
	tfs["right_leg"] = v;

    v[0] = 0.0885; //0.1;
    v[1] = 0.00228;
    v[2] = -0.0;
    v[3] = -0.02945;
	tfs["left_leg"] = v;

    v[0] = 0.056; //0.1;
    v[1] = 0.003545;
    v[2] = 0.00075; 
    v[3] =-0.03097;
	tfs["right_knee"] = v;

    v[0] = 0.056; //0.1;
    v[1] = 0.003545;
    v[2] = -0.00075;
    v[3] = -0.03097;
	tfs["left_knee"] = v;

    v[0] = 0.013; //0.2;
    v[1] = 0.0;
    v[2] = 0.0; 
    v[3] =-0.016;
	tfs["right_ankle"] = v;

    v[0] = 0.013; //0.2;
    v[1] = 0.0;
    v[2] = -0.0;
    v[3] = -0.016;
	tfs["left_ankle"] = v;

    v[0] = 0.077; //0.2;
    v[1] = 0.0;
    v[2] = -0.005;
    v[3] = -0.001;
	tfs["right_foot"] = v;

    v[0] = 0.077; //0.2;
    v[1] = 0.0;
    v[2] = 0.005;
    v[3] = -0.001;
	tfs["left_foot"] = v;

}

void CoMUtilities::init_robovie_x_mass()
{
    robovie_x_mass = 0.0;
    map< string, vector<double> >::iterator it;
    for (it = tfs.begin(); it != tfs.end(); ++it)
      robovie_x_mass += it->second[0];
    cout << "robovie_x_mass is " << robovie_x_mass << endl; // should be around 1.3kg
}

void CoMUtilities::showComX()
{
    int resolution = 20;
    int com_pos = (int)((com(0) + support_polygon_x)*resolution/(2*support_polygon_x));
    
    if (com(0) < -support_polygon_x)
        cout << " CoM | ";
    else
        cout << " . | ";

    for(int i = 0; i < resolution; i++)
        if( i == com_pos && !(com(0) < -support_polygon_x))
            cout << "CoM ";
        else
            cout << ". ";
    
    if(com(0) > support_polygon_x)
        cout << "| CoM"<< endl;
    else
        cout << "| ." << endl;
}

void CoMUtilities::visualize_links_coms(pcl::PointCloud<pcl::PointXYZ>::Ptr coms)
{

    p->removePointCloud ("coms");
    p->addPointCloud (coms, "coms");
    p->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 3, "coms");

    p->spinOnce();

}

void CoMUtilities::visualize_com(pcl::PointCloud<pcl::PointXYZ>::Ptr com_pcl)
{
    p->removePointCloud ("com_pcl");
    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> single_color(com_pcl, 0, 255, 0);
    p->addPointCloud<pcl::PointXYZ> (com_pcl, single_color, "com_pcl");
    p->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 3, "com_pcl");

    p->spinOnce();
}

int CoMUtilities::is_com_in_support_polygon()
{
    //if (com(2) < 0)
    //  return 2;

    if (abs(com(0)) > support_polygon_x /*|| abs(com(1)) > support_polygon_y*/)
      return 1;
    else
      return 0;
}

Eigen::Vector3f CoMUtilities::tf_source2target(const string & source_frame, const string & target_frame)
{
    tf::StampedTransform transform;
    Eigen::Vector3f transform_xyz;
    
    try
    {
      //listener->waitForTransform(source_frame, target_frame, ros::Time(0), ros::Duration(10.0) ); // TODO: capire se ruba le tf ad Rviz!
      listener->lookupTransform(source_frame, target_frame, ros::Time(0), transform);

      transform_xyz(0) = transform.getOrigin().x();
      transform_xyz(1) = transform.getOrigin().y();
      transform_xyz(2) = transform.getOrigin().z();
    }
    catch (tf::TransformException ex)
    {
      ROS_ERROR("%s", ex.what());
    }

    return transform_xyz;
}

void CoMUtilities::calculate_com()
{
    cout << "Computing CoM..." << endl;
    string source_frame = "right_foot"; // which is on the ground

    pcl::PointCloud<pcl::PointXYZ>::Ptr coms(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::PointCloud<pcl::PointXYZ>::Ptr com_pcl(new pcl::PointCloud<pcl::PointXYZ>);

    com(0) = com(1) = com(2) = 0.0;

    map< string, vector<double> >::iterator it;
    for (it = tfs.begin(); it != tfs.end(); ++it)
    {
      string link_name = it->first;
      Eigen::Vector3f offset_com_ref(it->second[1], it->second[2], it->second[3]);
      Eigen::Vector3f local_ref_frame = tf_source2target(source_frame, link_name);
      double link_mass = it->second[0];

      //cout << "link " << link_name << endl;
      //cout << "offset_com_ref is " << offset_com_ref(0) << " " << offset_com_ref(1) << " " << offset_com_ref(2) << endl;
      //cout << "local_ref_frame is " << local_ref_frame(0) << " " << local_ref_frame(1) << " " << local_ref_frame(2) << endl;
      //cout << "link_mass is " << link_mass << endl;
      
      Eigen::Vector3f link_com = local_ref_frame + offset_com_ref;
      com += link_com * link_mass;

      pcl::PointXYZ link_com_xyz(link_com(0),  link_com(1), link_com(2));
      coms->points.push_back(link_com_xyz);
    }

    // x red y green z blue
    com = com / robovie_x_mass;

    cout << "CoM is " << com(0) << " " << com(1) << " " << com(2) << endl;

    showComX();

    if (pcl_com_visualization)
        visualize_links_coms(coms);

    pcl::PointXYZ com_xyz(com(0), com(1), com(2));
    com_pcl->points.push_back(com_xyz);
    pcl::PointXYZ com_xy(com(0), com(1), 0.0);
    com_pcl->points.push_back(com_xy);

    if (pcl_com_visualization && com(2) < 5) // scarto i valori troppo sballati altrimenti si impalla il visualizzatore, da fare anche per il robot
      visualize_com(com_pcl);
}
